#!/usr/bin/env python3

from getopt import getopt
from glob import glob
import os
import os.path
import re
import subprocess
import sys

def parse_key_value_pair(line):
    m = re.match(r'(\w+)=([0-9a-f]+)', line)
    key = m.group(1)
    val = int(m.group(2), 16)
    return key, val

def warn(message):
    print('WARNING:', message, file=sys.stderr)

def main():
    run_all = False
    build_dir = os.getenv('BUILD_DIR')
    test_dir = 'tests'
    if build_dir == '':
        build_dir ='/tmp/myarm2'
    flag_c = False

    opts, args = getopt(sys.argv[1:], 'ac')
    for k, v in opts:
        if k == '-a':
            run_all = True
        elif k == '-c':
            flag_c = True

    if run_all:
        test_names = [path[6:-2] for path in glob(f'{test_dir}/*.s')]
    else:
        test_names = args

    if not flag_c:
        make_proc = subprocess.run(['make', '-C', test_dir] + [name + '.bin' for name in test_names])
        if make_proc.returncode != 0:
            sys.exit(1)

    fail_num = 0

    for test_name in test_names:

        vcd_file_path = f'{build_dir}/{test_name}.vcd'
        uart_out_file_path = f'{build_dir}/{test_name}.out'
        if flag_c:
            image_file_path = f'{test_name}/{test_name}.bin'
        else:
            image_file_path = f'{test_dir}/{test_name}.bin'

        sim_args = [
            f'{build_dir}/verilated/Vsystem',
            '-t', vcd_file_path,
            '-u', uart_out_file_path,
        ]

        conf_file_path = f'{test_dir}/{test_name}.conf'
        if os.path.isfile(conf_file_path):
            with open(conf_file_path) as f:
                for line in f:
                    match = re.match(r'(\w+)=(.*)$', line)
                    k, v = match.group(1, 2)
                    if k == 'time_limit':
                        sim_args += ['-m', v]
                    else:
                        warn(f'unknown config option {k} in {conf_file_path}')

        mem_file_path = f'{test_dir}/{test_name}.mem'
        dump_file_path = None
        if os.path.isfile(mem_file_path):
            with open(mem_file_path) as f:
                dump_file_path = f'{build_dir}/{test_name}.dump'
                dump_start_s, dump_len_s, dump_data_s = f
                dump_start = int(dump_start_s[:-1], 16)
                dump_len = int(dump_len_s[:-1], 16)
                dump_data = bytes.fromhex(dump_data_s[:-1])
                sim_args += [
                    '-s', str(dump_start),
                    '-l', str(dump_len),
                    '-d', dump_file_path,
                ]

        sim_args.append(image_file_path)

        Vsystem = subprocess.run(sim_args, stdout=subprocess.PIPE)
        Vsystem_out = str(Vsystem.stdout, 'utf8')
        d = {}
        for line in Vsystem_out.splitlines():
            m = re.match(r'(\w+)=([0-9a-f]+)$', line)
            key = m.group(1)
            val = int(m.group(2), 16)
            d[key] = val

        passed = True

        reg_file_path = f'{test_dir}/{test_name}.reg'
        if os.path.isfile(reg_file_path):
            with open(reg_file_path) as f:
                for line in f:
                    key, val = parse_key_value_pair(line)
                    if d.get(key) != val:
                        sys.stderr.write(f'{key}: expected {val:#x} observed {d.get(key):#x}\n')
                        passed = False

        # write gtkwave file
        gtkwave_file_path = f'{build_dir}/{test_name}.gtkw'
        with open(gtkwave_file_path, 'w') as f:
            vcd_file_size = os.stat(vcd_file_path).st_size
            f.write(f'''\
[dumpfile] "{vcd_file_path}"
[dumpfile_size] {vcd_file_size}
[savefile] "{gtkwave_file_path}"
TOP.system.cpu.pipectl__i_de1_valid
TOP.system.cpu.pipectl__i_de2_valid
TOP.system.cpu.pipectl__i_ex1_valid
TOP.system.cpu.pipectl__i_ex2_valid
TOP.system.cpu.pipectl__i_ex3_valid
TOP.system.cpu.pipectl__i_ex4_valid
TOP.system.cpu.pipectl__i_wb_valid
-
TOP.system.cpu.pc__we
TOP.system.cpu.pc__br
TOP.system.cpu.pc__q[29:0]
-
TOP.system.cpu.htrans[1:0]
TOP.system.cpu.hwrite
TOP.system.cpu.haddr[31:0]
TOP.system.cpu.hsize[2:0]
TOP.system.cpu.hwdata[31:0]
TOP.system.cpu.hready
TOP.system.cpu.hresp
TOP.system.cpu.hrdata[31:0]
-
TOP.system.cpu.rf__i_we0
TOP.system.cpu.rf__i_wa0[3:0]
TOP.system.cpu.rf__i_wa0_alt_mode
TOP.system.cpu.rf__i_d0[31:0]
-
TOP.system.cpu.rf__i_we1
TOP.system.cpu.rf__i_wa1[3:0]
TOP.system.cpu.rf__i_wa1_alt_mode
TOP.system.cpu.rf__i_d1[31:0]
-
TOP.system.cpu.cpsr__i_nzcv_we
TOP.system.cpu.cpsr__i_nzcv_d[3:0]
TOP.system.cpu.cpsr__i_aif_we[2:0]
TOP.system.cpu.cpsr__i_aif_d[2:0]
TOP.system.cpu.cpsr__i_mode_we
TOP.system.cpu.cpsr__i_mode_d[4:0]
TOP.system.cpu.ex1__i_cpsr[31:0]
''')

        if dump_file_path is not None:
            with open(dump_file_path, 'rb') as f:
                observed = f.read()
                expected = dump_data
                if observed != expected:
                    sys.stderr.write(f'memory contents mismatch at {dump_start:#x}\n')
                    sys.stderr.write('expected:\n')
                    print(' '.join('%02x' % b for b in expected), file=sys.stderr)
                    sys.stderr.write('observed:\n')
                    print(' '.join('%02x' % b for b in observed), file=sys.stderr)
                    passed = False

        if passed:
            prompt = r'[32mPASS[m'
        else:
            prompt = r'[31mFAIL[m'
            fail_num += 1

        sys.stderr.write(f'{prompt} {test_name}\n')

    if fail_num == 0:
        print('All tests passed.')
    else:
        print(f'{fail_num} test(s) failed.')
        sys.exit(1)

if __name__ == '__main__':
    main()
